#!/usr/bin/python
#Needs 2 files:
#1. Tab-delimited file of annotation (XX fields)
#Example:
#CDS     SGRA_0024       19169   19669   +       -       SGRA_0024       -       -       -       hypothetical protein    
#CDS     SGRA_0025       19758   20810   +       fbaB    SGRA_0025       4.1.2.13        COG1830 DhnA-type fructose-1,6-bisphosphate aldolase and related enzymes        -       fructose-bisphosphate aldolase  K01623
#CDS     SGRA_0026       21279   22622   +       trkA    SGRA_0026       -       COG0569 K+ transport systems, NAD-binding component     -       potassium transporter peripheral membrane component     K03499
#CDS     SGRA_0027       22630   24132   +       trkH    SGRA_0027       -       COG0168 Trk-type K+ transport systems, membrane components      -       trk system potassium uptake protein     K03498
#CDS     SGRA_0028       24403   24810   -       -       SGRA_0028       -       COG1328 Oxygen-sensitive ribonucleoside-triphosphate reductase  -       hypothetical protein    
#CDS     SGRA_0029       24843   25202   -       -       SGRA_0029       -       -       -       hypothetical protein    
#CDS     SGRA_0030       25233   27389   +       -       SGRA_0030       3.4.14.5        COG1506 Dipeptidyl aminopeptidases/acylaminoacyl-peptidases     -       peptidase S9B dipeptidylpeptidase IV domain protein     K01278

#2. Tab-delimited file of new definitions, votes, length checks, and conserved domains. Example:
#SGRA_0083       117.94  Largest vote = 1        hypothetical protein CHU_1675   -
#SGRA_0085       19.463  Largest vote = 4        lipoprotein     SNase
#SGRA_0086       40.146  Largest vote = 3        cell surface receptor IPT/TIG domain protein    TIG
#SGRA_0096       106.46  Largest vote = 6        conserved hypothetical protein  -
#SGRA_0101       132.33  Largest vote = 25       putative amino acid transporter         ANF_receptor

import sys
import re
from Bio import SeqIO

tab_file = sys.argv[1]
tf = open(tab_file, "rU")
lines = tf.readlines()

update_file = sys.argv[2]
uf = open(update_file, "rU")
ulines = uf.readlines()

m1 = re.compile('conserved hypothetical protein.*', re.IGNORECASE)
m2 = re.compile('hypothetical protein \w+.*', re.IGNORECASE)
m3 = re.compile('PREDICTED.*', re.IGNORECASE)

#new_defs = {}
to_check_locus = []
to_check_lengths = []
to_check_votes = []
to_check_desc = []

for line in ulines:
    l = line.split('\t')
    locus_tag = l[0]
    lengths = l[1]
    if l[2] == "No BLAST hit":
        votes = 0
    else:
        v = l[2].split(' ')
        votes = v[3]
    desc = l[3]
    domains = l[4].rstrip()
    #new_defs[locus_tag] = desc
    to_check_locus.append(locus_tag)
    to_check_lengths.append(lengths)
    to_check_votes.append(votes)
    to_check_desc.append(desc)

for line in lines:
    l = line.split('\t')
    feat = l[0]
    locus_tag = l[1]
    start = l[2]
    stop = l[3]
    frame = l[4]
    gene = l[5]
    protein_id = l[6]
    ec_num = l[7]
    note = l[8]
    bound_moiety = l[9]
    desc = l[10].rstrip()
    ko = l[11].rstrip()
    if locus_tag in to_check_locus:
        #assign stuff
        index = to_check_locus.index(locus_tag)
        #updated_def = new_defs[locus_tag]
        length_to_check = to_check_lengths[index]
        votes = to_check_votes[index]
        updated_def = to_check_desc[index]
        if 80 < length_to_check < 120: #check if length is within range
            if votes == 0:
                updated_def = "hypothetical protein"
            elif votes == 1:
                updated_def = "hypothetical protein"
            else:
                #do the regex matching to re-assign definitions
                if m1.match(updated_def):
                    updated_def = "conserved hypothetical protein"
                elif m2.match(updated_def):
                    updated_def = "conserved hypothetical protein"
                elif m3.match(updated_def):
                    updated_def = "hypothetical protein"
        else: #assign as hypotheticals if length is out of range
            if votes == 0:
                updated_def = "hypothetical protein"
            elif votes == 1:
                updated_def = "hypothetical protein"
            else:
                #do the regex matching to re-assign definitions
                if m1.match(updated_def):
                    updated_def = "conserved hypothetical protein"
                elif m2.match(updated_def):
                    updated_def = "conserved hypothetical protein"
                elif m3.match(updated_def):
                    updated_def = "hypothetical protein"
        #print final lines
        print feat + "\t" + locus_tag + "\t" + start + "\t" + stop + "\t" + frame \
            + "\t" + gene + "\t" + protein_id + "\t" + ec_num + "\t" + note + "\t" \
            + bound_moiety + "\t" + updated_def + "\t" + ko
    else:
        print feat + "\t" + locus_tag + "\t" + start + "\t" + stop + "\t" + frame \
            + "\t" + gene + "\t" + protein_id + "\t" + ec_num + "\t" + note + "\t" \
            + bound_moiety + "\t" + desc + "\t" + ko

tf.close()
uf.close()